package com.zjj.lbw.ai.old.rag;

import org.springframework.ai.document.Document;
import org.springframework.ai.reader.TextReader;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.io.Resource;
import org.springframework.stereotype.Component;

import java.util.List;

/**
 * 读取文本
 */
@Component
public class DocumentService {

    @Value("classpath:meituan-qa.txt") // This is the text document to load
    private Resource resource;

    @Autowired
    private VectorStore vectorStore;

    public List<Document> loadText() {
        TextReader textReader = new TextReader(resource);
        textReader.getCustomMetadata().put("filename", "meituan-qa.txt");
        List<Document> documents = textReader.get();

        ZjjTextSplitter zjjTextSplitter = new ZjjTextSplitter();
//        TokenTextSplitter tokenTextSplitter = new TokenTextSplitter();
//        List<Document> list = tokenTextSplitter.apply(documents);
        List<Document> list = zjjTextSplitter.apply(documents);

        // 把问题存到元数据中
        list.forEach(document -> document.getMetadata().put("question", document.getContent().split("\\n")[0]));

        // 向量存储
        vectorStore.add(list);

        return list;
    }

    public List<Document> search(String message){
        List<Document> documents = vectorStore.similaritySearch(message);
        return documents;
    }

    public List<Document> metadataSearch(String message, String question) {
        return vectorStore.similaritySearch(
                SearchRequest
                        .query(message)
                        .withTopK(5)
                        .withSimilarityThreshold(0.9)
                        .withFilterExpression(String.format("question in ['%s']", question)));
    }
}
